package main

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/tmc/langchaingo/llms/ollama"
	"net/http"
	"ollama-gin/models"
)

func main() {
	engine := gin.Default()
	v1 := engine.Group("/api/v1")
	{
		v1.GET("/", TestPong)
		v1.POST("/generate", GenerateHandler)
	}
	err := engine.Run(":8080")
	if err != nil {
		return
	}
}
func TestPong(ctx *gin.Context) {
	ctx.JSON(http.StatusOK, gin.H{
		"message": "Hello World",
	})
}

// GenerateHandler 处理xxxxx
func GenerateHandler(ctx *gin.Context) {
	param := ctx.Param("prompt")
	req := &models.RequestData{}
	req.Prompt = param
	if err := ctx.Bind(req); err != nil {
		fmt.Println(err)
		ResError(err.Error(), ctx)
		return
	}
	ollamaURL := "http://localhost:11434"
	llm, err := ollama.New(ollama.WithModel("llama3"), ollama.WithServerURL(ollamaURL))
	if err != nil {
		ResError(err.Error(), ctx)
		return
	}
	call, err := llm.Call(context.Background(), req.Prompt)
	if err != nil {
		ResError(err.Error(), ctx)
		return
	}
	fmt.Println(call)
	ctx.JSON(http.StatusOK, gin.H{
		"Response": call,
	})
	return
}

func ResError(error string, ctx *gin.Context) {
	ctx.JSON(http.StatusBadRequest, gin.H{
		"error": error,
	})
}
